fix(aiguard): scope framework collision avoidance to request and response phases - #20461

Description

Jira: APPSEC-70286, APPSEC-70474, APPSEC-70480 and APPSEC-70282

AI Guard used a single all-or-nothing depth counter to stop a framework integration and a provider integration from evaluating the same call twice. While it was raised, provider listeners skipped both their request and their response evaluation. LangChain streaming raised it on iteration entry but never evaluated the streamed response, so with DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED=true a streamed LangChain response was scanned by nobody.

Phase-scoped claims

The counter is replaced by independent REQUEST and RESPONSE claims, so a framework only suppresses the provider checks it performs itself:

Provider .before listeners check REQUEST; provider .after listeners and the buffered streams check RESPONSE.

AI Guard owns its LangChain logic; the contrib stays product-agnostic

The LangChain contrib (ddtrace/contrib/internal/langchain/) carries no AI Guard state or naming. It only dispatches the product-agnostic langchain.{chatmodel,llm}.{generate,agenerate,stream}.before and .{generate,agenerate}.after events (AI Guard no longer listens to .stream.before; see below), from inside the LLM span, so a listener's block is still recorded on that span and the LLMObs output is kept. Against main the contrib diff is removals only, plus renaming the stream option aiguard_before_event to before_event: the AI Guard state dicts, the aiguard_*_event options, the stream handler's start_stream override and the .finally / .stream.started / .stream.finally events are gone. They existed only so AI Guard could release a claim in a different listener from the one that took it.

Instead, at langchain.patch AI Guard installs its own wrappers (outermost, around the contrib's) on BaseChatModel / BaseLLM:

• generate / agenerate: with aiguard_context(REQUEST, RESPONSE) around the call.
• _stream / _astream of each model class: the stream buffer below, which also takes the stream's claim. Installed at langchain.patch for every existing subclass and from an __init_subclass__ hook for classes defined later, so a path that reads _stream directly (stream_events(version="v3")) is covered on first use.

Every claim is taken and released in one frame, so no handle travels through the contrib and nothing can leave a claim behind: not a nested inner stream, not an early break, not a close from another task.

The claim covers only the read of the model's stream, so the caller's loop body and LangChain callbacks are never inside it. This fixes APPSEC-70480: direct OpenAI / Anthropic calls made while iterating a LangChain stream are evaluated. Previously the claim spanned the whole iteration and those calls skipped AI Guard.

Streamed LangChain responses are buffered around each model's _stream / _astream

With DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED=true, AI Guard wraps the _stream / _astream each model class defines. The whole model stream is read, the chunks are merged into the message LangChain would build, request + response are evaluated, and only then are the chunks replayed. Subclasses define these methods themselves, so the wrappers are installed per class (see above). The layer matters:

• Below LangChain's callbacks. stream(), astream_events(), LangGraph stream_mode="messages" and generate's internal streaming all hand each token to callbacks after reading _stream. No callback sees a token before the verdict, and astream_events() still emits its stream events because the model run is still open when the first chunk is replayed. Providers that report tokens to the run manager from inside _stream (LLMs, ChatOpenAI(streaming=True)) get a deferring run manager whose tokens are replayed after the verdict.
• Streaming inside invoke() (APPSEC-70474). When LangChain streams inside generate (a streaming callback handler, streaming=True, stream=True), the buffer evaluates the response before any token is reported. It records the payload it evaluated, and .generate.after skips that exact payload, so the response is evaluated once. The record is scoped to the generate call, so a cached response is still evaluated.
• Timeouts. langchain-openai wraps each provider read inside _astream in its own 120 s wait_for. The buffer drains _astream from outside, so every provider read keeps its own timeout.
• Blocks cannot be swallowed. The buffer raises the BaseException-based AIGuardAbortError. The provider's OpenAIAIGuardAbortError is an Exception, which with_fallbacks and retry policies caught and routed to an unevaluated fallback.
• No double tool-call evaluation. On a clean verdict the evaluated tool calls are recorded in the last chunk's response_metadata; chunk addition carries it into the aggregate, so the legacy AgentExecutor (which streams by default) and the langgraph hooks skip them.
• The buffer evaluates the request too. Before claiming and reading the model, it evaluates the request unless that model's own generate call is running (tracked per model, so a different model streamed inside a generate still gets its check). Without this, a direct _stream read (stream_events v3) had no request check while its claim made the provider skip its own, so with the flag off its request was evaluated by nobody. This makes the langchain.*.stream.before listeners redundant; they are removed, so stream(), astream() and direct reads take one path. A blocked stream() now raises on its first next() rather than at the stream() call, and LangChain's on_llm_error callback fires for the block.
• Native protocol events. langchain-core 1.4 chat models can implement _stream_chat_model_events / _astream_chat_model_events, which stream_events / astream_events v3 read instead of _stream. Those hooks get the same per-class buffer: the request is evaluated first, the events are drained and rebuilt into a message with LangChain's ChatModelStream, evaluated, then replayed, so on_stream_event callbacks see nothing before the verdict. Tool calls on this path are not marked as evaluated, so an agent hook may evaluate them again (one extra evaluation, not a missed one).
• The provider iterator is always closed. Sync and async, buffered or passthrough, the model's iterator is closed after a full drain, a failure or a cancellation, so class-based SDK streams are not left open.
• Patched classes stay collectable. The per-class record of wrapped methods is a WeakKeyDictionary, so model classes defined at runtime can be garbage-collected.
• Any provider. ChatAnthropic and every other chat model or LLM are covered, not only OpenAI. The buffer is the evaluator, so it does not step aside for an outer framework's claim: a LangChain stream inside a Strands tool is evaluated.
• Cost. One claim per stream instead of one per chunk, and the chunks are merged in one add_ai_message_chunks call instead of pairwise.

It is passthrough when the flag is off. Conversion errors fail open; only a block raises. A subclass _stream that calls super()._stream is buffered once: the buffer skips only a nested read of the model it is already reading, on every read of that model (so delegating after a first chunk is not buffered again), while the claim still covers only the first read. Another model streamed during that read (a router model's inner model) gets its own buffer.

Claims are shared objects, released by handle (APPSEC-70282)

An asyncio task works on a copy of its parent's Context, so a counter lowered from another task left the claiming task covered for good, and later OpenAI/Anthropic calls in it skipped AI Guard. Claims are now shared objects stored in a ContextVar: a release from any task is seen by all, never raises (the old ContextVar.reset raised ValueError into the framework's cleanup path), and never drops an unrelated claim.

Every claim is released by its handle, never by searching the current context for "the latest claim", which misses claims made in another task and takes an inner call's claim when calls nest (the tokenless helper is removed). LangChain claims and releases in one frame via aiguard_context. Strands, whose own hooks split claim and release, keeps the handle on invocation_state.

Smaller fixes

• The Anthropic .after event skips streamed requests: a with_raw_response stream reaches it with its body unread, and converting it raised and reported a converter error on every call.
• is_aiguard_context_active has an early return and a plain loop (it runs on every provider call).
• Contributor guide (.cursor/rules/ai-guard.mdc) documents which phase each framework claims and each provider check reads, never claiming a phase you do not evaluate, the AI Guard-owned LangChain wrappers, why the stream buffer sits at _stream, and release by handle (never add AI Guard state to a contrib).

Out of scope

Pre-existing issues raised in review are tracked separately in Jira and are not fixed here: REQUEST claimed when LangChain skips its request check, the Strands invocation-wide RESPONSE claim over Strands' own models, the Strands structured_output_async claim leak, the async raw-response parse() coroutine, async evaluation blocking the event loop, Strands Graph parallel nodes sharing one handle, and making the phase a required argument. The double request evaluation when stream() falls back to invoke() (APPSEC-70482) no longer happens here: with the .stream.before listeners removed, that path is evaluated once, by .generate.before.

Testing

• Context primitive (test_context.py): phase scoping; cross-task release (no raise, clears the claiming task, keeps unrelated claims); a generator closed from another task releases its claim.
• LangChain (test_langchain.py): streamed chat (sync, async) and LLM responses evaluated before delivery, with the evaluated content checked against the delivered chunks; a response block delivers zero chunks; with_fallbacks does not swallow a block; a model with no provider integration is covered; a streaming AgentExecutor (stream_runnable=True) evaluates its tool call once; streamed tool calls marked evaluated for the agent hooks; claims observed at the provider's HTTP request for streamed and non-streamed calls; the caller's loop body is unclaimed with and without the buffer (APPSEC-70480); a _generate returning with an inner stream still open leaves no claim; an astream() closed from another task leaves no claim; conversion errors fail open.
  New in this revision: callback tokens and on_llm_end arrive only after the verdict, and none arrive on a block; astream_events() on a bare model emits every stream event, after the verdict; a chain's astream_events() emits no model output on a block; invoke(stream=True) is evaluated once, before its tokens (APPSEC-70474); tokens a model reports from inside _stream / _astream wait for the verdict and are dropped on a block (chat sync, async and LLM); a super()._stream subclass is buffered once; one claim per stream, with and without the buffer; a stream under an outer framework claim is still evaluated; the one-call chunk merge matches chunk addition; an inner model streamed by a router model is buffered on its own, so its callbacks see nothing on a block; _stream read directly is buffered for classes defined before and after patching; unpatch removes the hook and the wrappers.
  Review follow-ups: a _stream read directly evaluates its request and blocks before the model is read (deny, abort); a cancelled buffered _astream closes the provider stream; a runtime-defined model class is garbage-collected after buffering; a stream() that falls back to invoke() (no _stream, disable_streaming=True) evaluates its request once (APPSEC-70482); the provider iterator is closed after a drain and after a failure; a subclass that delegates to super()._stream() after its first chunk is evaluated once; native protocol events (_stream_chat_model_events, via stream_events(version="v3")) are buffered, delivered on allow, blocked on deny, and a denied request blocks before the model is read. All of them fail on the previous code.
• Provider buffer (test_streaming.py): unchanged coverage of evaluate-then-replay and blocking.
• Anthropic: the after event leaves a streamed raw response untouched and reports no converter error.
• Strands: cross-task release through the real hooks.

Cassettes: 2 files in tests/cassettes/openai/ for the OpenAI-patched streaming requests (stream_options.include_usage changes the body); response bodies reuse a recorded stream plus a usage chunk.

Local results on the final code:

Skips are the langchain-version-gated agent tests, plus the astream_events(version="v2") tests on langchain-core 0.1 and the native protocol-event tests before langchain-core 1.4. tests/llmobs/test_base_stream_handler.py: 94 passed locally (the .stream.finally pairing test was removed with the event); its async tests need pytest-asyncio, which the local llmobs venv lacks (CI covers them).

Lint clean: fmt, typing, spelling; ast-grep's no-rst-double-backticks warnings in the touched packages dropped from 58 to 49, none on added lines.

Risks

Behaviour change for streamed LangChain responses when DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED=true. They are now buffered and evaluated where previously they streamed through unevaluated: higher time-to-first-token, for the caller and for LangChain callbacks (astream_events(), LangGraph messages streaming, models streaming inside invoke()), and one additional evaluation per streamed LangChain call. The flag is off by default, so nobody is affected without opting in. Called out in the release note.

The LangChain contrib change is removals only: it no longer dispatches the AI Guard-only .finally / .stream.started / .stream.finally events (AI Guard was their only listener) and no longer carries AI Guard options. _context.py is private; its no-argument defaults are kept for existing callers and will be removed in a follow-up.

AI Guard wraps LangChain model methods directly, as it already does for agent plan and langgraph ToolNode. The _stream / _astream wrappers and the __init_subclass__ hook are installed at langchain.patch and removed on unpatch. Core ExecutionContexts were considered instead, but a context entered at stream start and exited at finalize would change the caller's current core context for the whole loop, and would stay stuck in the starting task when an async stream is finalized in another task.

Additional Notes

The OpenAI buffered-stream wrappers are installed at openai.patch time only if the flag is already on (existing behaviour); the OpenAI-patched LangChain fixture enables it before patching. The LangChain model-stream buffer checks the flag per call.

🤖 Generated with Claude Code